
import torch
import torch.nn as nn
import torch.nn.functional as F

# ***************************************************************************************
# brief: 这个类是一个 自定义损失函数模块，类似于 nn.CrossEntropyLoss
#        label smoothing 的作用就是将 one-hot vector -> 更加柔和的 vector
#        比如： [0, 0, 1, 0, 0] -> [0.025, 0.025, 0.9, 0.025, 0.025]，这样模型就不会执着于把
#        0.9 变成 1 了，这样模型就不会把所有的注意力都集中在一个“绝对正确”的标签上，
#        而是学会在一定程度上容忍不确定性。
#        具体的过程就是 将 one-hot vector 中 正确类别的概率设置为 confidence = 1 - ε = 0.9
#        其余类别平均分配剩下的 ε = 0.1
#
# @f (forward): 输入 x 是 一个 batch 中所有的 tokens 的预测 tensor，
#        它的 shape 是 (batch_size * seq_len, vocab_size), 也就是说每一行都是一个预测 tokens tensor
#        一共有 batch_size * seq_len 个 token 预测向量
#        其实就是把所有的 预测 tensor 按照列的方向全部堆叠起来形成了这个 tensor
#        一般 x 都是 model 的实际输出 再 经过 log_softmax，比如 x 长这样：
#        tensor([[-0.6515, -2.1515, -1.6515, -2.3515, -2.5515],   # token 1 的预测概率
#                [-3.1693, -2.0693, -2.8693, -0.8693, -2.2693],   # token 2 的预测概率
#                [-0.8915, -2.1915, -2.2915, -2.3915, -2.0915],   # token 3 的预测概率
#                [-3.1704, -2.1704, -1.1704, -2.6704, -2.4704],   # token 4 的预测概率
#                [-2.3793, -2.4793, -2.1793, -2.0793, -1.4793],   # token 5 的预测概率
#                [-1.7995, -1.9995, -2.8995, -2.6995, -2.5995]])  # token 6 的预测概率
#
#        输入 target 是 一个 batch 中理论输出结果的合集，shape 是 (batch_size * seq_len)
#        比如 [21, 503, 7, 38, 193, ...], 就表明 第一个 tokens 的目标输出应该是
#        词表中的第 21 个词，第二个 tokens 的目标输出应该是 词表中的 第 503 个词 ......
#
#        然后此函数会根据 x 和 taget 构建起一个 label smoothing 的 label，并使用 KLDivLoss
#        来计算最终的 loss
# =======================================================================================
class LabelSmoothing(nn.Module):

    def __init__(
            self,
            size,
            padding_idx: int,   # padding_idx 就是 padding 符号在 词表 中的代号
            smoothing=0.0,
    ):
        super(LabelSmoothing, self).__init__()
        # 损失函数使用 KL散度，度量模型输出和目标分布之间的距离
        self.criterion = nn.KLDivLoss(reduction="sum")
        # 忽略掉一些标签
        self.padding_idx = padding_idx
        # confidence 是给正确标签的 分值，其余的 label 平分 smoothing
        self.confidence = 1.0 - smoothing
        # 其他类平分 smoothing
        self.smoothing = smoothing
        self.size = size

        self.true_dist = None

    # forward 函数本质上是 1. 把 label 平滑,  2. 计算实际输出和 label smoothing 之间的 loss
    # x 是模型的 对数概率（log_softmax 输出）
    def forward(
            self,
            # 每一 **行** 都对应了一个 seq 预测出的 tokens vector
            x: torch.Tensor,         # shape:(batch_size * seq_len, vocab_size)
            target: torch.Tensor,    # shape:(batch_size * seq_len)
    ):
        assert x.size(1) == self.size  # 理论上应该都是 batch_size * seq_len

        # 构建 label smoothing 的 taget >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        # 借用 x 的形状
        true_dist = x.detach().clone()

        # 将所有类别的概率初始化为 smoothing / (size - 2)
        # size-2 是因为 padding 标签和正确标签排除掉了
        true_dist.fill_(self.smoothing / (self.size - 2))

        # 对正确类别位置赋值 confidence（如 0.9）
        # tensor.scatter_ (dim, index, value)
        # 在 dim=1 维度上，将 target.detach().unsqueeze(1) 的位置的值全部 替换为  self.confidence
        true_dist.scatter_(1, target.detach().unsqueeze(1), self.confidence)

        # 把 self.padding_idx 那一列的全部设置为 0
        true_dist[:, self.padding_idx] = 0

        # 创建一个掩码索引列表，标出 target 中哪些位置是 <pad>
        # 因为 taget 也是 batch 起来的，所以有的时候会兼顾其他数据 导致 这个 位置是 padding_idx
        mask = torch.nonzero(target.detach() == self.padding_idx)

        if mask.dim() > 0:
            # dim=0 对行进行操作，将 mask.squeeze() 的位置的数值全部设置为 0
            true_dist.index_fill_(0, mask.squeeze(), 0.0)

        self.true_dist = true_dist
        # 构建 label smoothing 的 taget <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

        # self.criterion -> KLDivLoss
        return self.criterion(x, true_dist.detach().clone())


# *****************************************************************************************
# brief: 输入 model 的原始输入 和 目标输出， 输出所有 tokens 的 总loss 和 平均 loss(用于梯度下降)
#
# @input(x): model 的输出(经过了 linear + softmax)，shape 是 (batch_size, seq_len, vocab_size)
# @input(y): 模型的目标输出，shape 是 (batch_size, seq_len)
#
# @ret: (total_loss, avg_loss)
# =========================================================================================
def computeLoss(
        x: torch.Tensor,
        y: torch.Tensor,
        criterion: LabelSmoothing,
        norm
) -> tuple[torch.Tensor, torch.Tensor]:
    sloss = (
            criterion(
                # 将 x 的 shape
                # 从 (batch_size, seq_len, vocab_size) -> (batch_size * seq_len, vocab_size)
                # batch_size * seq_len 就是所有 tokens 的数量
                # 就相当于把所有的 预测的 tokens 向量 全部拉平
                # 把 y 的 shaper 从 (batch_size, seq_len) -> (batch_size * seq_len)
                x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)
            )
            / norm
        # norm 一般是 tokens, 来表示 每个 tokens 的平均 loss
    )
    # sloss.detach() * norm 是总的 loss
    # sloss 才是真正参与梯度下降的 loss
    return sloss.detach() * norm, sloss


if __name__ == "__main__":
    # # 假设词表大小为 5
    # vocab_size = 5
    # padding_idx = 0
    # smoothing = 0.1
    #
    # # 构造模型输出 (log_softmax 前)
    # logits = torch.tensor([
    #     [[2.0, 0.2, 0.1, -1.0, 0.5],  # token 1
    #      [1.0, 0.0, 0.1, -1.5, 0.4],  # token 2
    #      [0.5, 0.1, 0.0, 0.2, 0.3]],  # token 3
    #
    #     [[-0.5, 1.0, 0.2, 0.0, 0.8],  # token 4
    #      [0.2, 0.5, -0.1, 0.3, 0.0],  # token 5
    #      [-0.5, 0.0, 1.0, 0.5, -0.3]]  # token 6
    # ])  # shape: (2, 3, 5)
    #
    # # 目标输出 (token 对应的词表索引)
    # targets = torch.tensor([
    #     [1, 2, 3],
    #     [1, 0, 2]  # 假设 token 5 是 <pad>，padding_idx=0
    # ])  # shape: (2, 3)
    #
    # # 构造 LabelSmoothing 模块
    # label_smoothing = LabelSmoothing(size=vocab_size, padding_idx=padding_idx, smoothing=smoothing)
    #
    # # 对 logits 做 log_softmax，因为 KLDivLoss 要求输入是 log-probability
    # log_probs = F.log_softmax(logits, dim=-1)  # shape: (2, 3, 5)
    #
    # # 计算 loss
    # norm = torch.sum(targets != padding_idx)  # 有效 token 数量（排除 padding）
    # total_loss, avg_loss = computeLoss(log_probs, targets, label_smoothing, norm)
    #
    # print(f"有效 token 数量 norm = {norm.item()}")
    # print(f"总 loss = {total_loss.item():.4f}")
    # print(f"平均 loss (供反向传播用) = {avg_loss.item():.4f}")

    # 假设 vocab_size = 5, target token 是索引为 2
    vocab_size = 5
    target_idx = 2

    target = torch.tensor([target_idx])
    padding_idx = 0

    # 创建 label smoothing loss
    criterion = LabelSmoothing(size=vocab_size, padding_idx=padding_idx, smoothing=0.1)


    # helper：打印每个 case 的 loss
    def print_loss(logits, desc):
        log_probs = F.log_softmax(logits, dim=-1)  # 必须是 log 概率
        loss = criterion(log_probs, target)
        print(f"{desc:25s} loss = {loss.item():.4f}")


    # 各种 logits 情况下的 loss 比较
    print_loss(torch.tensor([[0.0, 0.0, 5.0, 0.0, 0.0]]), "✅ 极度接近 target")
    print_loss(torch.tensor([[0.0, 0.0, 2.0, 0.0, 0.0]]), "✅ 接近 target")
    print_loss(torch.tensor([[0.0, 0.0, 0.5, 0.0, 0.0]]), "🟡 略接近 target")
    print_loss(torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0]]), "🟡 全部一样（乱猜）")
    print_loss(torch.tensor([[5.0, 0.0, 0.0, 0.0, 0.0]]), "❌ 预测成完全错误")

# ✅ 极度接近 target             loss = 0.0917
# ✅ 接近 target                loss = 0.1977
# 🟡 略接近 target              loss = 0.8465
# 🟡 全部一样（乱猜）                loss = 1.1745
# ❌ 预测成完全错误                 loss = 4.5917

